# -*- coding: utf-8 -*-
import json
import uuid
from typing import List, Optional

from loguru import logger
from pydantic import BaseModel
from fastapi import Request
from zai import ZhipuAiClient
from sse_starlette import EventSourceResponse

# ==================== 配置参数 ====================
# GLM API 配置
GLM_API_KEY = "fabc0f1f271448f9a0b7207ae47f77c5.o8ITXRoUG1N1dTEe"


# ==================== 基础工具函数 ====================
# 调用GLM API
async def call_glm_api(messages, json_mode=False):
    """调用GLM API"""
    try:
        client = ZhipuAiClient(api_key=GLM_API_KEY)
        
        response = client.chat.completions.create(
            model="glm-4.6",
            messages=messages,
            thinking={
                "type": "enabled"
            },
            response_format={
                "type": "json_object"
            } if json_mode else None,
            stream=True
        )
        
        for chunk in response:
            if not chunk.choices:
                continue

            content = chunk.choices[0].delta.content

            if content and content.strip():
                # logger.debug(f"chunk: {text}")
                yield content
            else:
                yield ''
    except Exception as e:
        logger.error(f"GLM调用失败: {e}")
        import traceback
        logger.error(traceback.format_exc())
        yield ''


async def chat_iter(query: str, dialog_id: str, origin_query: str, token: str):
    prompt = """
    你是一个专业的港口行业安全事故案例分析助手，旨在帮助港口从业人员快速检索、深入理解并从事故案例中学习，以提升整体作业安全水平。你需要通过多维度和模糊查询精准匹配事故案例，并以结构化的方式清晰呈现事故的基本信息、多层次原因分析（直接、间接、管理）、关键违规点、具体的防范措施及相关法规。此外，你应具备深度分析能力，能够针对用户的追问进行因果剖析、责任判定和风险评估，并根据培训、应急、排查等不同场景，调整回答的重点和策略。在交互中，请保持专业且易懂的语言风格，针对性回答用户提问，始终将安全教育和风险预防放在首位，并确保所有案例信息真实、客观且具有教育意义。
    用户的问题为：
    """
    try:
        messages = [
            {"role": "system", "content": prompt},
            {"role": "user", "content": query}
        ]

        res = ''
        async for i in call_glm_api(messages):
            if i is None:
                break
            res += i if i else ''
            if i:
                yield i

        final_response = json.dumps({"response": res}, ensure_ascii=False)
        yield final_response
        logger.info(f"dialog_id: {str(dialog_id)}, question: {origin_query}, answer: {res}")

    except Exception as e:
        import traceback
        traceback.print_exc()
        error_response = json.dumps({
            "event": "error",
            "data": json.dumps({
                "code": 500,
                "data": None,
                "msg": "回答生成失败"
            }, ensure_ascii=False)
        }, ensure_ascii=False)
        yield error_response


class ChatEntity(BaseModel):
    query: str
    dialog_id: Optional[uuid.UUID | str] = None
    token: str
    kb_ids: Optional[List[uuid.UUID]] = None


async def accident_chat(request: Request, payload: ChatEntity):
    """港口行业安全事故案例分析助手"""
    return EventSourceResponse(chat_iter(query=payload.query, dialog_id=payload.dialog_id, origin_query=payload.query, token=payload.token))


